import time
import path
from tqdm import tqdm

import gxl_ai_utils.thread.my_thread
import path
from gxl_ai_utils.utils import utils_file
from gxl_wv_mos import get_wvmos

model = get_wvmos(False)


def get_mos_score_for_single_wav(wav_path):
    return model.calculate_one(wav_path)


def get_mvmos_for_scp_file(wav_scp_file_path: str, target_scp_file_path: str):
    utils_file.makedir_for_file(target_scp_file_path)
    dic_data = utils_file.load_dic_from_scp(wav_scp_file_path)
    res_dic = {}
    num = 0
    print('开始得到得到wav.scp文件对应音频的mvmos, 共有{}条数据'.format(len(dic_data)))
    for k, v in tqdm(dic_data.items(), total=len(dic_data)):
        output = get_mos_score_for_single_wav(v)
        res_dic[k] = output
    utils_file.write_dic_to_scp(res_dic, target_scp_file_path)


def get_mvmos_for_scp_file_multi_thread(wav_scp_file_path: str, target_scp_file_path: str):
    utils_file.makedir_for_file(target_scp_file_path)
    dic_data = utils_file.load_dic_from_scp(wav_scp_file_path)
    res_dic = {}
    dic_data = {k: v for k, v in list(dic_data.items())[:100]}
    print('开始得到得到wav.scp文件对应音频的mvmos, 共有{}条数据'.format(len(dic_data)))

    def handle(dic_o):
        for k, v in tqdm(dic_o.items(), total=len(dic_o)):
            output = get_mos_score_for_single_wav(v)
            res_dic[k] = output

    thread_obj = gxl_ai_utils.thread.my_thread.MyThreadPool()
    num_thread = 4
    split = len(dic_data) // num_thread
    for i in range(0, len(dic_data), split):
        now_dic = {k: v for k, v in list(dic_data.items())[i:i + split]}
        thread_obj.add_thread(handle, [now_dic])
    thread_obj.start()
    utils_file.write_dic_to_scp(res_dic, target_scp_file_path)


if __name__ == '__main__':
    now = time.time()
    get_mvmos_for_scp_file_multi_thread("./output/wav.scp", "./output/mvmos.scp")
    print(f'time cost: {time.time() - now}')
